#pragma once
#include <vector>
#include <string>
#include <opencv2/opencv.hpp>
#include "onnxruntime_cxx_api.h"

#include "IInfer.h"

class InferenceOnnx : public IInfer
{
public:
    InferenceOnnx() {};
    ~InferenceOnnx();

    void setTaskType(task_type type) {};
    std::vector<Detection> runInfer(cv::Mat& src);
    bool loadModel(std::string modelPath);
    bool loadLabels(std::string labelFilePath);


    void setSrcImage(cv::Mat& src) { srcImage_ = src; }
    void setInferParams(float modelConfidence, float modelScore, float modelNMS)
    { 
        modelConfidenceThreshold_ = modelConfidence;
        modelScoreThreshold_ = modelScore;
        modelNMSThreshold_ = modelNMS;
    }

    std::string getProInfo() { return processInfo_; };

protected:
    void postProcess_v8(float *data, float factor);
    void postProcess_v5(float* data, float factor);
    // classesFile
    void parseTextFile(std::wstring filePath);

    void release();

private:
    cv::Mat srcImage_;
    Ort::Session *session_ = nullptr;
    Ort::SessionOptions* session_options_ = nullptr;;
    Ort::Env* env_ = nullptr;

	std::vector<std::string> labels_;

    float xFactor_ = 1.0;
    float yFactor_ = 1.0;

    int input_w = 0;
    int input_h = 0;

    int output_h = 0;
    int output_w = 0;

    std::vector<std::string> inputNodeNames_;
    std::vector<std::string> outputNodeNames_;

    std::vector<Detection> detections_;
    std::vector<std::string> timeTestRecord_;   // 记录测试时间

    float modelConfidenceThreshold_ = 0.25;
    float modelScoreThreshold_ = 0.25;
    float modelNMSThreshold_ = 0.45;

    std::string modelPath_;
    std::string processInfo_;
};

